from bs4 import BeautifulSoup
import re

def fill_missing_cells_with_dash(html_content):
    """将HTML表格中<tr>内缺失的单元格填充为——"""
    soup = BeautifulSoup(html_content, 'html.parser')
    
    # 查找表格
    table = soup.find('table')
    if not table:
        return str(soup)
    
    # 查找表头行，确定列数
    max_cols = 0
    header_rows = []
    
    # 检查<thead>中的行
    thead = table.find('thead')
    if thead:
        for row in thead.find_all('tr'):
            header_rows.append(row)
            cols = len(row.find_all(['td', 'th']))
            # 考虑colspan
            colspan_total = 0
            for cell in row.find_all(['td', 'th']):
                colspan = cell.get('colspan')
                if colspan:
                    colspan_total += int(colspan) - 1
            cols += colspan_total
            max_cols = max(max_cols, cols)
    
    # 检查<tbody>中的行
    tbody = table.find('tbody')
    body_rows = tbody.find_all('tr') if tbody else table.find_all('tr')
    
    # 计算最大列数
    for row in body_rows:
        # 跳过已经在header_rows中的行
        if row in header_rows:
            continue
        cols = len(row.find_all(['td', 'th']))
        # 考虑colspan
        colspan_total = 0
        for cell in row.find_all(['td', 'th']):
            colspan = cell.get('colspan')
            if colspan:
                colspan_total += int(colspan) - 1
        cols += colspan_total
        max_cols = max(max_cols, cols)
    
    # 为每行填充缺失的单元格
    for row in header_rows + body_rows:
        # 跳过重复的行
        if row in header_rows and row in body_rows:
            continue
            
        # 计算当前行的列数（考虑colspan）
        current_cols = len(row.find_all(['td', 'th']))
        colspan_total = 0
        for cell in row.find_all(['td', 'th']):
            colspan = cell.get('colspan')
            if colspan:
                colspan_total += int(colspan) - 1
        current_cols += colspan_total
        
        # 如果当前行的列数小于最大列数，添加缺失的单元格
        missing_cols = max_cols - current_cols
        if missing_cols > 0:
            for _ in range(missing_cols):
                # 创建新的<td>元素并添加到行末尾
                new_cell = soup.new_tag('td')
                new_cell.string = "——"
                row.append(new_cell)
    
    return str(soup)

def html_table_to_markdown(html_content):
    """将HTML表格转换为Markdown格式"""
    # 首先填充缺失的单元格
    filled_html = fill_missing_cells_with_dash(html_content)
    
    # 使用BeautifulSoup解析填充后的HTML
    soup = BeautifulSoup(filled_html, 'html.parser')
    
    # 查找表格
    table = soup.find('table')
    if not table:
        return ""
    
    # 初始化Markdown表格
    markdown_table = []
    
    # 处理表头
    thead = table.find('thead')
    if thead:
        for row in thead.find_all('tr'):
            header_row = []
            for cell in row.find_all(['th', 'td']):
                # 移除多余的空白字符和换行符
                cell_text = cell.get_text(strip=True).replace('\n', '<br>')
                header_row.append(cell_text)
            if header_row:
                markdown_table.append('| ' + ' | '.join(header_row) + ' |')
                # 添加分隔行
                separator_row = '| ' + ' | '.join(['---'] * len(header_row)) + ' |'
                markdown_table.append(separator_row)
    
    # 处理表体
    tbody = table.find('tbody')
    rows = tbody.find_all('tr') if tbody else table.find_all('tr')
    
    for row in rows:
        # 跳过表头行（如果在tbody之外）
        if thead and row in thead.find_all('tr'):
            continue
            
        data_row = []
        for cell in row.find_all(['td', 'th']):
            # 移除多余的空白字符和换行符
            cell_text = cell.get_text(strip=True).replace('\n', '<br>')
            data_row.append(cell_text)
        if data_row:
            markdown_table.append('| ' + ' | '.join(data_row) + ' |')
    
    return '\n'.join(markdown_table)

# HTML table content (provided by the user)
html_table = """
<table><thead><tr><th rowspan="2">介质</th><th rowspan="2">法兰公称压力<br>(MPa)</th><th rowspan="2">工作温度<br>(°C)</th><th rowspan="2">法兰型式</th><th rowspan="2">密封面</th><th colspan="2">垫片</th><th rowspan="2">备注</th></tr><tr><th>型式</th><th>材料</th></tr></thead><tbody><tr><td rowspan="6">油品、油气、溶<br>剂:石油化工<br>原料及产品<sup>b</sup></td><td rowspan="2">≤1.6</td><td>≤200</td><td>甲、乙型平焊</td><td>平(凹凸)面</td><td>耐油垫、四氟垫</td><td>耐油橡胶石棉板、聚四氟乙烯板</td><td rowspan="6">当介质为易燃、易<br>爆、有毒或强渗透<br>性时,应采用凹凸<br>面法兰</td></tr><tr><td>201~250</td><td>长颈对焊</td><td>平(凹凸)面</td><td>缠绕垫、金属包垫、<br>柔性石墨复合垫</td><td>0Cr13(0Cr18Ni9)钢带+特制石棉(石<br>墨)、铁皮(铝皮)+特制石棉、石墨+金属<br>骨架(0Cr13、0Cr18Ni9等)</td></tr><tr><td rowspan="2">2.5</td><td>≤200</td><td>乙型平焊</td><td>平(凹凸)面</td><td>耐油垫、缠绕垫、金属包垫、<br>柔性石墨复合垫</td><td>耐油橡胶石棉板<br>0Cr13(0Cr18Ni9)钢带+特制石棉(石<br>墨)、铁皮(铝皮)+特制石棉、石墨+金属<br>骨架(0Cr13、0Cr18Ni9等)</td></tr><tr><td>201~250</td><td>长颈对焊</td><td>平(凹凸)面</td><td>缠绕垫、金属包垫、<br>柔性石墨复合垫</td><td>0Cr13(0Cr18Ni9)钢带+特制石棉(石<br>墨)、铁皮(铝皮)+特制石棉、石墨+金属<br>骨架(0Cr13、0Cr18Ni9等)</td></tr><tr><td rowspan="2">4.0</td><td>≤40</td><td>长颈对焊</td><td>凹凸面</td><td>缠绕垫、柔性石墨<br>复合垫</td><td>0Cr13(0Cr18Ni9)钢带+特制石棉(石<br>墨)、石墨+金属骨架(0Cr13、0Cr18Ni9<br>等)</td></tr><tr><td>41~450</td><td>长颈对焊</td><td>凹凸面</td><td>缠绕垫、金属包垫、<br>柔性石墨复合垫</td><td>0Cr13(0Cr18Ni9)钢带+石墨带、铁皮<br>(铝皮)+特制石棉、石墨+金属骨架<br>(0Cr13、0Cr18Ni9等)</td></tr><tr><td>6.4</td><td>≤450</td><td>长颈对焊</td><td>凹凸面<br>梯形槽</td><td>缠绕垫、金属包垫<br>金属环垫</td><td>0Cr13(0Cr18Ni9)钢带+石墨带、铁皮<br>(0Cr13)+特制石棉<br>0Cr13、0Cr18Ni9、10</td></tr></tbody></table>
"""

# 包含缺失单元格的测试HTML表格
test_html_with_missing_cells = """
<table>
  <thead>
    <tr>
      <th>姓名</th>
      <th>年龄</th>
      <th>职业</th>
      <th>城市</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>张三</td>
      <td>25</td>
      <td>工程师</td>
      <td>北京</td>
    </tr>
    <tr>
      <td>李四</td>
      <td>26</td>
      <td>设计师</td>
      <!-- 缺失城市单元格 -->
    </tr>
    <tr>
      <td>王五</td>
      <td>30</td>
      <!-- 缺失职业和城市单元格 -->
    </tr>
    <tr>
      <td>赵六</td>
      <!-- 缺失年龄、职业和城市单元格 -->
    </tr>
  </tbody>
</table>
"""

# 示例用法
if __name__ == "__main__":
    # 转换原始表格
    markdown_result = html_table_to_markdown(html_table)
    print("原始表格转换结果:")
    print(markdown_result)
    
    # 保存到文件
    with open("test_table_md.md", "w", encoding="utf-8") as f:
        f.write(markdown_result)
    
    print("\n" + "="*50 + "\n")
    
    # 转换包含缺失单元格的测试表格
    test_markdown_result = html_table_to_markdown(test_html_with_missing_cells)
    print("包含缺失单元格的测试表格转换结果:")
    print(test_markdown_result)
    
    # 保存到文件
    with open("test_result.md", "w", encoding="utf-8") as f:
        f.write(test_markdown_result)
    
    print("\n转换完成，结果已保存到文件。")